#include"spsolver.h"
#include<iostream>
#include<fstream>
#include<cmath>
#include<numeric>
#include<iomanip>
using std::cout;
using std::endl;
using std::ifstream;
using std::accumulate;
using std::setw;


void SPSolver::setup()
{
  SystemHF.setup();
  SystemWS.setup();
  //setup Channels
  Channels.clear();
  ifstream fin("channel.dat");
  int l,jj,tz;
  while(true)
    {
      fin>>l>>jj>>tz;
      if(tz==0) break;
      int point=0;
      Channels.push_back( SPChannel(l,jj,tz) );
    }
}


double SPSolver::hf_r(int i,double r) const
{
  return hfsolver.hf_r(i,r);
}

// KDataType SPSolver::getHFk1k2(KDataType k1,KDataType k2,int group)
// {
//   KDataType valk1k2=0;
//   int bra=pSystem_Tz->Groups[group][0];
    
//   int jja=pSystem_Tz->Orbitals[bra].jj;
//   int tza=pSystem_Tz->Orbitals[bra].tz;
//   int la=pSystem_Tz->Orbitals[bra].l;

//   for(int i=0;i<hfsolver.HFOrbitals.size();i++)
//     {
//       if(hfsolver.HFOrbitals[i].state) continue;//skip particles
//       int jjh=hfsolver.HFOrbitals[i].jj;
//       int Jmin=abs(jja - jjh)/2;
//       int Jmax=(jja + jjh)/2;
//       if(Jmin>Jmax) continue;
//       int GroupIndex=hfsolver.HFOrbitals[i].GroupIndex;
//       int dim=pSystem_Tz->Groups[GroupIndex].size();

//       //k1 \dot k2
//       if(hfsolver.HFOrbitals[i].tz!=tza) continue;
//       if((hfsolver.HFOrbitals[i].l+la+1)%2) continue;//parity
//       for(int J=Jmin;J<=Jmax;J++)
// 	{
// 	  valk1k2-=(2*J+1)*(jjh+1) * pow(gsl_sf_coupling_3j(jja,2,jjh,1,0,-1),2) * gsl_sf_coupling_6j(jja,jjh,J*2,jja,jjh,2) * k1 * k2 * hf_k(i,k1) * hf_k(i,k2);
// 	}

//     }
//   return valk1k2;
// }

double SPSolver::getVrrfromHF(int channel,double r1,double r2)
{
  int group;
  if(findChannelInHO(channel,group))
    {
      return getVrrfromHF(r1,r2,group);
    }
  else
    return 0;
}

double SPSolver::getVrrfromHF(double r1,double r2,int group)
{
  double Vrr=0;
  double lengthb=pSystem_Tz->lengthb;
  const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
  int l=HOOrbital.l;
  int jj=HOOrbital.jj;
  int tz=HOOrbital.tz;
  int dim=pSystem_Tz->Groups[group].size();
  for(int i=0;i<dim;i++)
    for(int j=0;j<dim;j++)
      {
	const System_Tz_HO::OrbitalType & HOOrbital1=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][i] ];
	const System_Tz_HO::OrbitalType & HOOrbital2=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][j] ];
	int n1=HOOrbital1.n;
	int n2=HOOrbital2.n;
	int bra=pSystem_Tz->Groups[group][i];
	int ket=pSystem_Tz->Groups[group][j];
	
        Vrr+=SHFDiagram(bra,ket)*ho_r(n1,l,lengthb,r1)*ho_r(n2,l,lengthb,r2);
      }
  return Vrr;
}
void SPSolver::printHFVrr(int channel,int numpoints,double step)
{
  int totalGroups=pSystem_Tz->Groups.size();

  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  cout<<l<<"\t"<<jj<<"\t"<<tz<<endl;
  for(int i=0;i<numpoints;i++)
    for(int j=0;j<numpoints;j++)
      {
	double r1=i*step;
	double r2=j*step;
	cout<<r1<<"\t"<<r2<<"\t"<<getVrrfromHF(channel,r1,r2)<<endl;
      }
}

bool SPSolver::findChannelInWS(const int & channel,int &group) const
{
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  int totalGroups=SystemWS.Groups.size();
  group=0;//find the group of l,jj,tz
  for(group=0;group<totalGroups;group++)
    {
      const System_Tz_WS::OrbitalType & WSOrbital=SystemWS.Orbitals[ SystemWS.Groups[group][0] ];
      if(WSOrbital.l==l && WSOrbital.jj==jj && WSOrbital.tz==tz)
	break;
    }
  if(group>=totalGroups)
    {
      std::cout<<"warning: channel: "<<l<<","<<jj<<","<<tz<<" not exist in WS basis, skipped\n";
      return false;
    }
  return true;
}

bool SPSolver::findChannelInHO(const int & channel,int &group) const
{
  int l=Channels[channel].l;
  int jj=Channels[channel].jj;
  int tz=Channels[channel].tz;
  int totalGroups=pSystem_Tz->Groups.size();
  group=0;//find the group of l,jj,tz
  for(group=0;group<totalGroups;group++)
    {
      const System_Tz_HO::OrbitalType & HOOrbital=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[group][0] ];
      if(HOOrbital.l==l && HOOrbital.jj==jj && HOOrbital.tz==tz)
	break;
    }
  if(group>=totalGroups)
    {
      std::cout<<"warning: channel: "<<l<<","<<jj<<","<<tz<<" not exist in HO basis(so in HF potential), skipped\n";      
      return false;
    }
  return true;
}
void SPSolver::setVMat(int channel,KMat &h)
{
  int dim;
  int group;
  if(findChannelInWS(channel,group))
    {
      dim=SystemWS.Groups[group].size();
    }
  else
    dim=0;

  h.resize(dim,dim);
  h.setZero();

  #pragma omp parallel for
  for(int i=0;i<dim;i++)
    for(int j=i;j<dim;j++)
      {
	int bra=SystemWS.Groups[group][i];
	int ket=SystemWS.Groups[group][j];
	h(i,j)=SystemWS.get1B(realA,bra,ket);//kinetic
	const System_Tz_WS::OrbitalType & WSOrb_bra=SystemWS.Orbitals[bra];
	const System_Tz_WS::OrbitalType & WSOrb_ket=SystemWS.Orbitals[ket];      

	for(int nbra=0;nbra<WSOrb_bra.r.size();nbra++)
	  for(int nket=0;nket<WSOrb_ket.r.size();nket++)
	    {
	      double r1=WSOrb_bra.r[nbra];
	      double w1=WSOrb_bra.w[nbra];
	      double r2=WSOrb_ket.r[nket];
	      double w2=WSOrb_ket.w[nket];
	      h(i,j)+=WSOrb_bra.wavefun[nbra]*w1 *r1*getVrrfromHF(channel,r1,r2)*r2* w2*WSOrb_ket.wavefun[nket];
	    }	
	h(j,i)=h(i,j);
      }
}

double SPSolver::SHFDiagram(int bra,int ket) const
{
  if (iter==0)
    return hfsolver.SHFDiagram(bra,ket);
  
  int jj=pSystem_Tz->Orbitals[bra].jj;
  KDataType val2B=0;
  for(int i=0;i<=FermiSurface;i++)
    {
      int Jmin=abs(jj - Orbitals[i].jj)/2;
      int Jmax=(jj + Orbitals[i].jj)/2;


      int GroupIndex=Orbitals[i].GroupIndex;
      int dim=pSystem_Tz->Groups[GroupIndex].size();

      for(int k=0;k<dim;k++)
	for(int l=0;l<dim;l++)
	  {
	    for(int J=Jmin;J<=Jmax;J++)
	      {
		int brak=pSystem_Tz->Groups[GroupIndex][k];
		int ketl=pSystem_Tz->Groups[GroupIndex][l];
		val2B += hfsolver.getTwoBody(bra,brak,ket,ketl,J) * Orbitals[i].Coeff[k] * Orbitals[i].Coeff[l] * KDataType(2*J+1);
	      }
	  }
    }
  val2B/=(jj+1);
  return val2B.real();
}

void SPSolver::solve()
{
  Orbitals.clear();
  for(int channel=0;channel<Channels.size();channel++)
    {
      KMat h;      
      setVMat(channel,h);
      int dim=h.cols();
      if(dim==0) continue;
      EigenSolver es(h);
      for(int a=0;a<dim;a++)
  	{
  	  KDataType norm=sqrt( es.eigenvectors().col(a).transpose()*es.eigenvectors().col(a) );
  	  // if(complexd(100.,0) < es.eigenvalues()[a])
  	  //   continue;
  	  Orbitals.push_back( SP_R(Channels[channel].l,Channels[channel].jj,Channels[channel].tz,channel,es.eigenvalues()[a],es.eigenvectors().col(a)/norm) );
  	}
    }


  // for(int channel=0;channel<Channels.size();channel++)
  //   {
  //     int group;
  //     if(findChannelInWS(channel,group))
  // 	{
  // 	  int dim=SystemWS.Groups[group].size();
  // 	  int l=Channels[channel].l;
  // 	  int jj=Channels[channel].jj;
  // 	  int tz=Channels[channel].tz;
  // 	  for(int i=0;i<dim;i++)
  // 	    {
  // 	      const System_Tz_WS::OrbitalType &Orb=SystemWS.Orbitals[ SystemWS.Groups[group][i] ];
  // 	      complexd e=Orb.e;
  // 	      KVec vec(dim);
  // 	      vec=KVec::Zero(dim);
  // 	      vec(i)=1.;
  // 	      Orbitals.push_back(SP_R(l,jj,tz,channel,e,vec));
  // 	      Orbitals[Orbitals.size()-1].type=Orb.type;
  // 	    }
  // 	}
  //   }
  

  //  setType();
  setupOverlap(Orbitals.size());//set overlap between HO. and gammow basis, used for interaction basis transform

  //the rest partial wave
  int totalGroups=SystemHF.Groups.size();
  for(int group=0;group<totalGroups;group++)
    {
      const System_Tz_HF::OrbitalType & Orb=SystemHF.Orbitals[ SystemHF.Groups[group][0] ];
      int l=Orb.l;
      int jj=Orb.jj;
      int tz=Orb.tz;
      int i=0;
      for(i=0;i<Channels.size();i++)
	{
	  if( Channels[i].l==l && Channels[i].jj==jj && Channels[i].tz==tz )
	    break;
	}
      if(i>=Channels.size())
	{
	  for(int k=0;k<SystemHF.Groups[group].size();k++)
	    {
	      const System_Tz_HF::OrbitalType & HFOrb=SystemHF.Orbitals[ SystemHF.Groups[group][k] ];
	      KDataType e=HFOrb.e;
	      int dim=HFOrb.Coeff.size();
	      SP_R sp(l,jj,tz,-1,e,KVec());
	      sp.type=-1;
	      sp.GroupIndex=HFOrb.GroupIndex;
	      sp.Coeff.resize(dim);	      
	      for(int n=0;n<dim;n++)
		sp.Coeff[n]=HFOrb.Coeff[n];
	      Orbitals.push_back(sp);
	    }
	}
    }
  stable_sort( Orbitals.begin(),Orbitals.end());

  //below for binding energy
  Energy=0;
  FermiSurface=SystemHF.FermiSurface;
  for(int i=0;i<=FermiSurface;i++)
    {
      KDataType t=get1B(i,i);
      Energy+=0.5*(Orbitals[i].e+t)*(Orbitals[i].jj+1.);
    }
}


void SPSolver::iterSolve()
{
  double sigma=0,eps=1e-4;
  const int maxIter=300;
  iter=0;
  solve();
  int totalOrbitals=Orbitals.size();
  iter++;
  do
    {
      vector<SP_R> lastOrbitals=Orbitals;
      solve();
      for(int i=0;i<totalOrbitals;i++)
  	{
  	  sigma+=abs(lastOrbitals[i].e-Orbitals[i].e);
  	}
      sigma/=totalOrbitals;
      ++iter;      
    }
  while(sigma>eps && iter<maxIter);
  
}


void SPSolver::printOrbitals(int num,int channel)
{
  if(num==-1) num=Orbitals.size();
  int printnum=0;
  for(int i=0;i<Orbitals.size();i++)
    {
      int chan=Orbitals[i].channel;
      if(channel==-1||channel==chan)
	{
	  //	  cout<<"=============================\n";
	  cout<<setw(4)<<Orbitals[i].l<<"\t"<<setw(4)<<Orbitals[i].jj<<"\t"<<setw(4)<<Orbitals[i].tz<<"\t"<<setw(10)<<Orbitals[i].e.real()<<"\t"<<setw(10)<<Orbitals[i].e.imag()<<"\t"<<setw(10)<<std::arg(Orbitals[i].e)*180/Pi<<"\t"<<abs(InOverOut(i))<<endl;//<<"\t"<<Orbitals[i].type<<endl;
	  // cout<<"-----------------------------\n";
	  // for(int j=0;j<Channels[chan].numPoints;j++)
	  //   {
	  //     cout<<Channels[chan].k[j].real()<<"\t"<<Channels[chan].wk[j].real()<<"\t"<<Orbitals[i].wavefun(j).real()<<endl;
	  //   }
	  // cout<<"=============================\n\n\n";
	  ++printnum;
	  if(printnum>=num) break;
	}
    }
}

void SPSolver::printOrb_r(int orb) const
{
  int group;
  findChannelInWS(Orbitals[orb].channel,group);
  const vector<double> & r=SystemWS.Orbitals[ SystemWS.Groups[group][0] ].r;

  vector<complexd> u(r.size(),0);
  for(int i=0;i<Orbitals[orb].wavefun.size();i++)
    {
      const vector<complexd> & wavefun=SystemWS.Orbitals[ SystemWS.Groups[group][i] ].wavefun;      
      for(int j=0;j<r.size();j++)
	{
	  u[j]+=wavefun[j]*Orbitals[orb].wavefun[i];
	}
    }
  for(int j=0;j<r.size();j++)
    {
      std::cout<<r[j]<<"\t"<<u[j].real()<<"\t"<<u[j].imag()<<"\t"<<hf_r(0,r[j])<<std::endl;
    }
}

complexd SPSolver::InOverOut(int orb) const
{
  int group;
  findChannelInWS(Orbitals[orb].channel,group);
  const vector<double> & vr=SystemWS.Orbitals[ SystemWS.Groups[group][0] ].r;
  int last=vr.size()-1;
  double r=vr[last];
  complexd u(0),up(0);
  for(int i=0;i<Orbitals[orb].wavefun.size();i++)
    {
      const vector<complexd> & wavefun=SystemWS.Orbitals[ SystemWS.Groups[group][i] ].wavefun;
      const vector<complexd> & ups=SystemWS.Orbitals[ SystemWS.Groups[group][i] ].up;
      u+=Orbitals[orb].wavefun[i]*wavefun[last];
      up+=Orbitals[orb].wavefun[i]*ups[last];
    }

  complexd k=sqrt(Orbitals[orb].e*2.*mc2/hbar_c/hbar_c/(1-1./realA));
  if(Orbitals[orb].e.real()<0) k=-k;
  int l=Orbitals[orb].l;
  complexd z=k*r;
  complexd jn,yn,djn,dyn;
  csphbessel_jy(l,z,jn,yn,djn,dyn);
  complexd ci(0,1.);
  complexd out=(-yn+ci*jn)*z;
  complexd in=(-yn-ci*jn)*z;
  complexd outp=(-dyn+ci*djn)*k*z+(-yn+ci*jn)*k;
  complexd inp=(-dyn-ci*djn)*k*z+(-yn-ci*jn)*k;
  //   /  out    in  \   / a \     / u \
  //  |               | |     | = |     |
  //   \  outp   inp /   \ b /     \ up/

  complexd bOvera=(out*up-outp*u)/(inp*u-in*up);
  return bOvera;  
}

KDataType SPSolver::get1B(int bra,int ket) const
{
  return (1.-1./realA)*get1B_p2(bra,ket);
}
KDataType SPSolver::get1B_p2(int bra,int ket) const
{
  KDataType val=0;
  if(Orbitals[bra].isInAGroupWith(Orbitals[ket]) )
    {
      int channel=Orbitals[ket].channel;
      if(channel==-1)//HF basis
	{
	  for(int ibra=0;ibra<Orbitals[bra].Coeff.size();ibra++)
	    for(int iket=0;iket<Orbitals[ket].Coeff.size();iket++)
	      {
		int iibra=pSystem_Tz->Groups[Orbitals[bra].GroupIndex][ibra];
		int iiket=pSystem_Tz->Groups[Orbitals[ket].GroupIndex][iket];
		val+=pSystem_Tz->get1B_p2(iibra,iiket)*Orbitals[bra].Coeff[ibra]*Orbitals[ket].Coeff[iket];
	      }	  
	}
      else
	{
	  int group;
	  findChannelInWS(channel,group);
	  for(int j=0;j<Orbitals[bra].wavefun.size();j++)
	    {
	      for(int l=0;l<Orbitals[ket].wavefun.size();l++)
		{
		  int ibra=SystemWS.Groups[group][j];
		  int iket=SystemWS.Groups[group][l];
		  val+=Orbitals[bra].wavefun(j)*SystemWS.get1B_p2(ibra,iket)*Orbitals[ket].wavefun(l);
		}
	    }
	}
     }
  return val;
}

void SPSolver::setType()
{
}


void SPSolver::setupOverlap(int num)
{
  double lengthb=pSystem_Tz->lengthb;
  int totalGroups=pSystem_Tz->Groups.size();

  for(int i=0;i<num;i++)
    {
      int group;
      int channel=Orbitals[i].channel;
      findChannelInWS(channel,group);
      
      Orbitals[i].GroupIndex=SystemWS.Orbitals[ SystemWS.Groups[group][0] ].GroupIndex;
      int dimws=SystemWS.Groups[group].size();
      int dimho=pSystem_Tz->Groups[Orbitals[i].GroupIndex].size();

      Orbitals[i].Coeff.resize(dimho);
      for(int j=0;j<dimho;j++)//every HO. orbitals in a group
	{
	  KDataType temp=0;//overlap with the HO. orbital
	  for(int k=0;k<dimws;k++)
	    {
	      const System_Tz_WS::OrbitalType & WSOrbital=SystemWS.Orbitals[ SystemWS.Groups[group][k] ];	      
	      temp += Orbitals[i].wavefun(k)*WSOrbital.Coeff[j];
	    }
	  Orbitals[i].Coeff[j]=temp;
	  //	  std::cout<<temp<<std::endl;
	}
    }
}
